#include <iostream>
#include <cassert>
#include <fstream>
#include <sstream>
#include <map>
#include "dataset.h"

using namespace std;


void dataset::setCategories(map<string,int>& category)
{
   // From http://cseweb.ucsd.edu/users/elkan/tabulate.html
    category["normal"] = 0;
    category["apache2"] = 2;
    category["back"] = 2;
    category["buffer_overflow"] = 3;
    category["ftp_write"] = 4;
    category["guess_passwd"] = 4;
    category["httptunnel"] = 4;
    category["httptunnel"] = 3;
    category["imap"] = 4;
    category["ipsweep"] = 1;
    category["land"] = 2;
    category["loadmodule"] = 3;
    category["mailbomb"] = 2;
    category["mscan"] = 1;
    category["multihop"] = 4;
    category["multihop"] = 3; // note that this is a duplicate
    category["named"] = 4;
    category["neptune"] = 2;
    category["nmap"] = 1;
    category["perl"] = 3;
    category["phf"] = 4;
    category["pod"] = 2;
    category["portsweep"] = 1;
    category["processtable"] = 2;
    category["ps"] = 3;
    category["rootkit"] = 3;
    category["saint"] = 1;
    category["satan"] = 1;
    category["sendmail"] = 4;
    category["smurf"] = 2;
    category["snmpgetattack"] = 4;
    category["snmpguess"] = 4;
    category["sqlattack"] = 3;
    category["teardrop"] = 2;
    category["udpstorm"] = 2;
    category["warezmaster"] = 2;
    category["worm"] = 4;
    category["xlock"] = 4;
    category["xsnoop"] = 4;
    category["xterm"] = 3;
    
    // training attack types
    
    category["back"] = 2;
    category["buffer_overflow"] = 3;
    category["ftp_write"] = 4;
    category["guess_passwd"] = 4;
    category["imap"] = 4;
    category["ipsweep"] = 1;
    category["land"] = 2;
    category["loadmodule"] = 3;
    category["multihop"] = 4;
    category["neptune"] = 2;
    category["nmap"] = 1;
    category["perl"] = 3;
    category["phf"] = 4;
    category["pod"] = 2;
    category["portsweep"] = 1;
    category["rootkit"] = 3;
    category["satan"] = 1;
    category["smurf"] = 2;
    category["spy"] = 4;
    category["teardrop"] = 2;
    category["warezclient"] = 4;
    category["warezmaster"] = 4;  // note that this contradicts the category above
}


dataset::dataset(string filename, int size)
{
   
    theRecords = new record*[size];
    nSize = 0;
  
    int nCount[] = {0,0,0,0,0};

    map<string,int> category;
    setCategories(category);

    ifstream inFile(filename.c_str());
    if (!inFile.is_open())
    assert(inFile.is_open());

    // PARSE INPUT FILE create new records
    string sLine;
    while (getline(inFile,sLine))
    {
        double duration = getNextNumericField(sLine);
        string protocolType = getNextStringField(sLine); 
        string service = getNextStringField(sLine); 
        string flag = getNextStringField(sLine);
        double src_bytes = getNextNumericField(sLine);
        double dst_bytes = getNextNumericField(sLine);
        double land = getNextNumericField(sLine);
        double wrongFragment = getNextNumericField(sLine);
        double urgent = getNextNumericField(sLine);
        double hot = getNextNumericField(sLine);
        double numFailedLogins = getNextNumericField(sLine);
        double loggedIn = getNextNumericField(sLine);
        double numCompromised = getNextNumericField(sLine);
        double rootShell = getNextNumericField(sLine);
        double suAttempted = getNextNumericField(sLine);
        double numRoot = getNextNumericField(sLine);
        double numFileCreations = getNextNumericField(sLine);
        double numShells = getNextNumericField(sLine);
        double numAccessFiles = getNextNumericField(sLine);
        double numOutboundCmds = getNextNumericField(sLine);
        double isHostLogin = getNextNumericField(sLine);
        double isGuestLogin = getNextNumericField(sLine);
        double count = getNextNumericField(sLine);
        double srv_count = getNextNumericField(sLine);
        double serrorRate = getNextNumericField(sLine);
        double srvSerrorRate = getNextNumericField(sLine);
        double rerrorRate = getNextNumericField(sLine);
        double srvRerrorRate = getNextNumericField(sLine);
        double sameSrvRate = getNextNumericField(sLine);
        double diffSrvRate = getNextNumericField(sLine);
        double srvDiffHostRate = getNextNumericField(sLine);
        double dstHostCount = getNextNumericField(sLine);
        double dstHostSrvCount = getNextNumericField(sLine);
        double dstHostSameSrvCount = getNextNumericField(sLine);
        double dstHostDiffSrvRate = getNextNumericField(sLine);
        double dstHostSameSrcPortRate = getNextNumericField(sLine);
        double dstHostSrcDiffHostRate = getNextNumericField(sLine);
        double dstHostSerrorRate = getNextNumericField(sLine);
        double dstHostSrvSerrorRate = getNextNumericField(sLine);
        double dstHostRerrorRate = getNextNumericField(sLine);
        double dstHostSrvRerrorRate = getNextNumericField(sLine);
    
        string attackName = sLine.substr(0,sLine.length()-1);
    
        assert(category.find(attackName) != category.end());
        int type = category[attackName];
    
        record* tmp = new record(type, duration, protocolType, service,
                                 flag, src_bytes, dst_bytes, land, wrongFragment,
                                 urgent, hot, numFailedLogins, loggedIn, numCompromised,
                                 rootShell, suAttempted, numRoot, numFileCreations,
                                 numShells, numAccessFiles, numOutboundCmds, isHostLogin,
                                 isGuestLogin, count, srv_count, serrorRate, 
                                 srvSerrorRate, rerrorRate, srvRerrorRate, 
                                 sameSrvRate, diffSrvRate, srvDiffHostRate, 
                                 dstHostCount, dstHostSrvCount, dstHostSameSrvCount,
                                 dstHostDiffSrvRate,
                                 dstHostSameSrcPortRate, dstHostSrcDiffHostRate, dstHostSerrorRate,
                                 dstHostSrvSerrorRate, dstHostRerrorRate, dstHostSrvRerrorRate);

        assert(nSize < size);
        theRecords[nSize] = tmp;
        nCount[type]++;
        nSize++;

    }

    cout << "Records Loaded:\n"
        << "  Normal:           " << nCount[0] << endl
        << "  Probe:            " << nCount[1] << endl
        << "  DoS:              " << nCount[2] << endl
        << "  User to Root:     " << nCount[3] << endl
        << "  Remote to Local:  " << nCount[4] << endl;
}

double dataset::getNextNumericField(string& sLine)
{
    double value;

    int nPos = sLine.find(",");
    string sField = sLine.substr(0,nPos); 
    sLine = sLine.substr(nPos+1);
    stringstream ss(sField);
    ss >> value;

    return value;
}

string dataset::getNextStringField(string& sLine)
{
    int nPos = sLine.find(",");

    string value = sLine.substr(0,nPos);
    sLine = sLine.substr(nPos+1);

    return value; 
}

string dataset::getFieldName(int nFld)
{
    if (nFld == 0)
        return "duration";
    else if (nFld == 1)
        return "protocolType";
    else if (nFld == 2)
        return "service";
    else if (nFld == 3)
        return "flag";
    else if (nFld == 4)
        return "src_bytes";
    else if (nFld == 5)
        return "dst_bytes";
    else if (nFld == 6)
        return "land";
    else if (nFld == 7)
        return "wrongFragment";
    else if (nFld == 8)
        return "urgent";
    else if (nFld == 9)
        return "hot";
    else if (nFld == 10)
        return "numFailedLogins";
    else if (nFld == 11)
        return "loggedIn";
    else if (nFld == 12)
        return "numCompromised";
    else if (nFld == 13)
        return "rootShell";
    else if (nFld == 14)
        return "suAttempted";
    else if (nFld == 15)
        return "numRoot";
    else if (nFld == 16)
        return "numFileCreations";
    else if (nFld == 17)
        return "numShells";
    else if (nFld == 18)
        return "numAccessFiles";
    else if (nFld == 19)
        return "numOutboundCmds";
    else if (nFld == 20)
        return "isHostLogin";
    else if (nFld == 21)
        return "isGuestLogin";
    else if (nFld == 22)
        return "count";
    else if (nFld == 23)
        return "srv_count";
    else if (nFld == 24)
        return "serrorRate";
    else if (nFld == 25)
        return "srvSerrorRate";
    else if (nFld == 26)
        return "rerrorRate";
    else if (nFld == 27)
        return "srvRerrorRate";
    else if (nFld == 28)
        return "sameSrvRate";
    else if (nFld == 29)
        return "diffSrvRate";
    else if (nFld == 30)
        return "srvDiffHostRate";
    else if (nFld == 31)
        return "dstHostCount";
    else if (nFld == 32)
        return "dstHostSrvCount";
    else if (nFld == 33)
        return "dstHostSameSrvCount";
    else if (nFld == 34)
        return "dstHostDiffSrvRate";
    else if (nFld == 35)
        return "dstHostSameSrcPortRate";
    else if (nFld == 36)
        return "dstHostSrcDiffHostRate";
    else if (nFld == 37)
        return "dstHostSerrorRate";
    else if (nFld == 38)
        return "dstHostSrvSerrorRate";
    else if (nFld == 39)
        return "dstHostRerrorRate";
    else if (nFld == 40)
        return "dstHostSrvRerrorRate";

    assert(false);

}

string dataset::getFieldValue(int nFld, int nValue)
{
    if (nFld == 1)
    {
        if (nValue == 0) return "tcp";
        if (nValue == 1) return "udp";
        if (nValue == 2) return "icmp";
    }
    else if (nFld == 2)
    {
        if (nValue == 0) return "auth";
        if (nValue == 1) return "bgp";
        if (nValue == 2) return "courier";
        if (nValue == 3) return "csnet_ns";
        if (nValue == 4) return "ctf";
        if (nValue == 5) return "daytime";
        if (nValue == 6) return "discard";
        if (nValue == 7) return "domain";
        if (nValue == 8) return "domain_u";
        if (nValue == 9) return "echo";
        if (nValue == 10) return "eco_i";
        if (nValue == 11) return "ecr_i";
        if (nValue == 12) return "efs";
        if (nValue == 13) return "exec";
        if (nValue == 14) return "finger";
        if (nValue == 15) return "ftp";
        if (nValue == 16) return "ftp_data";
        if (nValue == 17) return "gopher";
        if (nValue == 18) return "hostnames";
        if (nValue == 19) return "http";
        if (nValue == 20) return "http_443";
        if (nValue == 21) return "icmp";
        if (nValue == 22) return "imap4";
        if (nValue == 23) return "IRC";
        if (nValue == 24) return "iso_tsap";
        if (nValue == 25) return "klogin";
        if (nValue == 26) return "kshell";
        if (nValue == 27) return "ldap";
        if (nValue == 28) return "link";
        if (nValue == 29) return "login";
        if (nValue == 30) return "mtp";
        if (nValue == 31) return "name";
        if (nValue == 32) return "netbios_dgm";
        if (nValue == 33) return "netbios_ns";
        if (nValue == 34) return "netbios_ssn";
        if (nValue == 35) return "netstat";
        if (nValue == 36) return "nnsp";
        if (nValue == 37) return "nntp";
        if (nValue == 38) return "ntp_u";
        if (nValue == 39) return "other";
        if (nValue == 40) return "pm_dump";
        if (nValue == 41) return "pop_2";
        if (nValue == 42) return "pop_3";
        if (nValue == 43) return "printer";
        if (nValue == 44) return "private";
        if (nValue == 45) return "remote_job";
        if (nValue == 46) return "rje";
        if (nValue == 47) return "shell";
        if (nValue == 48) return "smtp";
        if (nValue == 49) return "sql_net";
        if (nValue == 50) return "ssh";
        if (nValue == 51) return "sunrpc";
        if (nValue == 52) return "supdup";
        if (nValue == 53) return "systat";
        if (nValue == 54) return "telnet";
        if (nValue == 55) return "tftp_u";
        if (nValue == 56) return "tim_i";
        if (nValue == 57) return "time";
        if (nValue == 58) return "urh_i";
        if (nValue == 59) return "urp_i";
        if (nValue == 60) return "uucp";
        if (nValue == 61) return "uucp_path";
        if (nValue == 62) return "vmnet";
        if (nValue == 63) return "whois";
        if (nValue == 64) return "X11";
        if (nValue == 65) return "Z39_50";
        if (nValue == 66) return "http_2784";
        if (nValue == 67) return "harvest";
        if (nValue == 68) return "aol";
        if (nValue == 69) return "http_8001";
        if (nValue == 70) return "red_i";
    }
    else if (nFld == 3)
    {
        if (nValue == 0) return "OTH";
        if (nValue == 1) return "REJ";
        if (nValue == 2) return "RSTO";
        if (nValue == 3) return "RSTOS0";
        if (nValue == 4) return "RSTR";
        if (nValue == 5) return "S0";
        if (nValue == 6) return "S1";
        if (nValue == 7) return "S2";
        if (nValue == 8) return "S3";
        if (nValue == 9) return "SF";
        if (nValue == 10) return "SH";
    }

    assert(false);
}

string dataset::getAttackName(int nValue)
{
    switch (nValue)
    {
    case 0:  return "normal"; 
    case 1:  return "probe"; 
    case 2:  return "dos"; 
    case 3:  return "u2r"; 
    case 4:  return "r2l"; 
    default:  assert(false);
    }
}

